package com.yf.limiter.impl;

import com.yf.limiter.RateLimiter;

import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 令牌桶限流算法
 * 核心思想：系统以固定速率生成令牌放入桶中，请求需要获取令牌才能执行，无令牌则拒绝或等待
 * 特点：既能限制平均速率，又能应对突发流量（桶内缓存的令牌），是最常用的限流算法之一
 */
public class TokenBucketRateLimiter<T> implements RateLimiter<T> {
    // 令牌桶容量（最大可缓存的令牌数）
    private final long bucketCapacity;
    // 令牌生成速率（单位：令牌/秒）
    private final double tokenGenerateRatePerSecond;
    // 当前桶内的令牌数（使用原子类保证线程安全）
    private final AtomicLong currentTokens = new AtomicLong(0);
    // 上次生成令牌的时间（毫秒时间戳）
    private volatile long lastGenerateTime;
    // 锁和条件变量（用于阻塞等待场景）
    private final ReentrantLock lock = new ReentrantLock();
    private final Condition hasToken = lock.newCondition();

    /**
     * 构造函数
     * @param bucketCapacity 令牌桶容量（最大缓存令牌数）
     * @param tokenGenerateRatePerSecond 令牌生成速率（令牌/秒）
     */
    public TokenBucketRateLimiter(long bucketCapacity, double tokenGenerateRatePerSecond) {
        this.bucketCapacity = bucketCapacity;
        this.tokenGenerateRatePerSecond = tokenGenerateRatePerSecond;
        this.lastGenerateTime = System.currentTimeMillis();
        this.currentTokens.set(0);
    }

    /**
     * 尝试获取1个令牌（非阻塞）
     * @return true：获取成功；false：获取失败（限流）
     */
    public boolean submit(T req) {
        return tryAcquire(1);
    }

    /**
     * 尝试获取指定数量的令牌（非阻塞）
     * @param tokens 需要获取的令牌数
     * @return true：获取成功；false：获取失败（限流）
     */
    public boolean tryAcquire(int tokens) {
        if (tokens <= 0) {
            throw new IllegalArgumentException("获取的令牌数必须大于0");
        }

        // 1. 先生成新的令牌
        generateTokens();

        // 2. 检查是否有足够的令牌
        long current = currentTokens.get();
        if (current < tokens) {
            return false; // 令牌不足，限流
        }

        // 3. 原子性减少令牌数（避免并发问题）
        return currentTokens.compareAndSet(current, current - tokens);
    }


    /**
     * 生成新的令牌（核心方法）
     */
    private void generateTokens() {
        long now = System.currentTimeMillis();
        long timeElapsedMs = now - lastGenerateTime;

        if (timeElapsedMs <= 0) {
            return; // 时间未推进，无需生成令牌
        }

        // 计算应该生成的令牌数 = 时间差（秒）* 生成速率
        double generateAmount = (timeElapsedMs / 1000.0) * tokenGenerateRatePerSecond;
        if (generateAmount <= 0) {
            return;
        }

        // 计算新的令牌数（不能超过桶容量）
        long newTokens = (long) Math.min(bucketCapacity, currentTokens.get() + Math.floor(generateAmount));
        currentTokens.set(newTokens);

        // 更新上次生成时间
        lastGenerateTime = now;

        // 唤醒等待令牌的线程（如果有）
        if (newTokens > 0) {
            lock.lock();
            try {
                hasToken.signalAll();
            } finally {
                lock.unlock();
            }
        }
    }

    // Getter方法（用于监控和调试）
    public long getBucketCapacity() {
        return bucketCapacity;
    }

    public double getTokenGenerateRatePerSecond() {
        return tokenGenerateRatePerSecond;
    }

    public long getCurrentTokens() {
        generateTokens(); // 先生成令牌再返回准确值
        return currentTokens.get();
    }
}
